{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9ca550b3",
   "metadata": {},
   "source": [
    "# debug"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1d6cb3af",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "import tvm\n",
    "from tvm import tir\n",
    "from tvm.relax.frontend import nn\n",
    "from tvm.relax.frontend.nn import op, spec\n",
    "from tvm.runtime import NDArray"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7e0874e",
   "metadata": {},
   "source": [
    "## `print`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d9ec7590",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_1225526/3381093976.py:3: shape = (10, 5), dtype = float32, data =\n",
      "[[0.73462075 0.01047635 0.31426018 0.8674046  0.12726694]\n",
      " [0.05869561 0.14252067 0.57591957 0.62454313 0.5909737 ]\n",
      " [0.04823518 0.8938983  0.89632434 0.15392512 0.18556875]\n",
      " [0.31064558 0.6404101  0.62769175 0.1271398  0.89232284]\n",
      " [0.16888297 0.4059453  0.70379585 0.77904415 0.84397256]\n",
      " [0.00197196 0.48467797 0.3471346  0.00262749 0.82645226]\n",
      " [0.28290933 0.51377964 0.67835134 0.04419875 0.07198626]\n",
      " [0.00947952 0.96474004 0.16065902 0.33178765 0.49370688]\n",
      " [0.6731469  0.5020207  0.44368893 0.4925446  0.45231378]\n",
      " [0.48918706 0.0983749  0.9165322  0.5107241  0.788272  ]]\n"
     ]
    }
   ],
   "source": [
    "class Layer(nn.Module):\n",
    "    def forward(self, x: nn.Tensor):  # pylint: disable=invalid-name\n",
    "        op.print_(x)\n",
    "        return x\n",
    "\n",
    "model = Layer().jit(\n",
    "    spec={\n",
    "        \"forward\": {\"x\": spec.Tensor([10, 5], dtype=\"float32\")},\n",
    "    },\n",
    "    debug=True,\n",
    ")\n",
    "x = torch.rand((10, 5), dtype=torch.float32)  # pylint: disable=invalid-name\n",
    "y = model[\"forward\"](x)  # pylint: disable=invalid-name\n",
    "assert isinstance(y, torch.Tensor)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4e7cb5e",
   "metadata": {},
   "source": [
    "## 调试函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "47a8b697",
   "metadata": {},
   "outputs": [
    {
     "ename": "AssertionError",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[4], line 32\u001b[0m\n\u001b[1;32m     22\u001b[0m model \u001b[38;5;241m=\u001b[39m Layer()\u001b[38;5;241m.\u001b[39mjit(\n\u001b[1;32m     23\u001b[0m     spec\u001b[38;5;241m=\u001b[39m{\n\u001b[1;32m     24\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mforward\u001b[39m\u001b[38;5;124m\"\u001b[39m: {\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     29\u001b[0m     debug\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     30\u001b[0m )\n\u001b[1;32m     31\u001b[0m x \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mrand((\u001b[38;5;241m10\u001b[39m, \u001b[38;5;241m5\u001b[39m), dtype\u001b[38;5;241m=\u001b[39mtorch\u001b[38;5;241m.\u001b[39mfloat32)  \u001b[38;5;66;03m# pylint: disable=invalid-name\u001b[39;00m\n\u001b[0;32m---> 32\u001b[0m y \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mforward\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m8\u001b[39;49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# pylint: disable=invalid-name\u001b[39;00m\n\u001b[1;32m     33\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(y, torch\u001b[38;5;241m.\u001b[39mTensor)\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/relax/frontend/nn/torch.py:78\u001b[0m, in \u001b[0;36mTorchModule.__getitem__.<locals>._closure\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m     71\u001b[0m args \u001b[38;5;241m=\u001b[39m [\n\u001b[1;32m     72\u001b[0m     _torch_to_tvm(arg_name, arg_spec, arg)\n\u001b[1;32m     73\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m arg_name, arg_spec, arg \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(\n\u001b[1;32m     74\u001b[0m         method_spec\u001b[38;5;241m.\u001b[39marg_names, method_spec\u001b[38;5;241m.\u001b[39marg_specs, args\n\u001b[1;32m     75\u001b[0m     )\n\u001b[1;32m     76\u001b[0m ]\n\u001b[1;32m     77\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39meffects \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m---> 78\u001b[0m     outputs, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39meffects \u001b[38;5;241m=\u001b[39m \u001b[43mmethod\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43meffects\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     79\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     80\u001b[0m     outputs \u001b[38;5;241m=\u001b[39m method(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mparams)\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/_ffi/_cython/packed_func.pxi:339\u001b[0m, in \u001b[0;36mtvm._ffi._cy3.core.PackedFuncBase.__call__\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/_ffi/_cython/packed_func.pxi:270\u001b[0m, in \u001b[0;36mtvm._ffi._cy3.core.FuncCall\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/_ffi/_cython/packed_func.pxi:259\u001b[0m, in \u001b[0;36mtvm._ffi._cy3.core.FuncCall3\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/_ffi/_cython/base.pxi:185\u001b[0m, in \u001b[0;36mtvm._ffi._cy3.core.CHECK_CALL\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/_ffi/base.py:468\u001b[0m, in \u001b[0;36mraise_last_ffi_error\u001b[0;34m()\u001b[0m\n\u001b[1;32m    462\u001b[0m \u001b[38;5;66;03m# The exception PyObject may contain a large amount of state,\u001b[39;00m\n\u001b[1;32m    463\u001b[0m \u001b[38;5;66;03m# including all stack frames that may be inspected in a later\u001b[39;00m\n\u001b[1;32m    464\u001b[0m \u001b[38;5;66;03m# PDB post-mortem.  Therefore, we must make sure to remove the\u001b[39;00m\n\u001b[1;32m    465\u001b[0m \u001b[38;5;66;03m# underlying PyObject* from the C++ side after we retrieve it.\u001b[39;00m\n\u001b[1;32m    466\u001b[0m _LIB\u001b[38;5;241m.\u001b[39mTVMDropLastPythonError()\n\u001b[0;32m--> 468\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m py_err\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/runtime/relax_vm/vm.cc:558\u001b[0m, in \u001b[0;36mtvm::runtime::relax_vm::VirtualMachineImpl::InvokeClosurePacked(tvm::runtime::ObjectRef const&, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)\u001b[0;34m()\u001b[0m\n\u001b[1;32m    556\u001b[0m   {\n\u001b[1;32m    557\u001b[0m     NVTXScopedRange scope(\"RelaxVM: \" + clo->func_name);\n\u001b[0;32m--> 558\u001b[0m     clo->impl.CallPacked(TVMArgs(values.data(), tcodes.data(), args.size() + 1), rv);\n\u001b[1;32m    559\u001b[0m   }\n\u001b[1;32m    560\u001b[0m }\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/runtime/relax_vm/vm.cc:632\u001b[0m, in \u001b[0;36moperator()\u001b[0;34m()\u001b[0m\n\u001b[1;32m    630\u001b[0m     inputs[i] = args[i + 1];\n\u001b[1;32m    631\u001b[0m   }\n\u001b[0;32m--> 632\u001b[0m   *rv = static_cast<VirtualMachineImpl*>(ctx_ptr)->InvokeBytecode(gf_idx, inputs);\n\u001b[1;32m    633\u001b[0m });\n\u001b[1;32m    634\u001b[0m return VMClosure(func_name, impl);\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/runtime/relax_vm/vm.cc:703\u001b[0m, in \u001b[0;36mtvm::runtime::relax_vm::VirtualMachineImpl::InvokeBytecode(long, std::vector<tvm::runtime::TVMRetValue, std::allocator<tvm::runtime::TVMRetValue> > const&)\u001b[0;34m()\u001b[0m\n\u001b[1;32m    701\u001b[0m   // set program counter\n\u001b[1;32m    702\u001b[0m   pc_ = gfunc.start_instr;\n\u001b[0;32m--> 703\u001b[0m   RunLoop();\n\u001b[1;32m    704\u001b[0m   return return_value_;\n\u001b[1;32m    705\u001b[0m }\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/runtime/relax_vm/vm.cc:828\u001b[0m, in \u001b[0;36mtvm::runtime::relax_vm::VirtualMachineImpl::RunLoop()\u001b[0;34m()\u001b[0m\n\u001b[1;32m    826\u001b[0m switch (instr.op) {\n\u001b[1;32m    827\u001b[0m   case Opcode::Call: {\n\u001b[0;32m--> 828\u001b[0m     this->RunInstrCall(curr_frame, instr);\n\u001b[1;32m    829\u001b[0m     break;\n\u001b[1;32m    830\u001b[0m   }\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/runtime/relax_vm/vm.cc:781\u001b[0m, in \u001b[0;36mtvm::runtime::relax_vm::VirtualMachineImpl::RunInstrCall(tvm::runtime::relax_vm::VMFrame*, tvm::runtime::relax_vm::Instruction)\u001b[0;34m()\u001b[0m\n\u001b[1;32m    779\u001b[0m \n\u001b[1;32m    780\u001b[0m   if (instrument_ == nullptr) {\n\u001b[0;32m--> 781\u001b[0m     this->InvokeClosurePacked(func_pool_[instr.func_idx], args, &ret);\n\u001b[1;32m    782\u001b[0m   } else {\n\u001b[1;32m    783\u001b[0m     // insert light-weight instrument callback\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/runtime/relax_vm/builtin.cc:497\u001b[0m, in \u001b[0;36moperator()\u001b[0;34m()\u001b[0m\n\u001b[1;32m    495\u001b[0m     }\n\u001b[1;32m    496\u001b[0m   }\n\u001b[0;32m--> 497\u001b[0m   debug_func->CallPacked(TVMArgs(call_args.data(), call_type_codes.data(), num_args + 1), rv);\n\u001b[1;32m    498\u001b[0m   *rv = io_effect;\n\u001b[1;32m    499\u001b[0m });\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/_ffi/_cython/packed_func.pxi:56\u001b[0m, in \u001b[0;36mtvm._ffi._cy3.core.tvm_callback\u001b[0;34m()\u001b[0m\n",
      "Cell \u001b[0;32mIn[4], line 10\u001b[0m, in \u001b[0;36m_debug\u001b[0;34m(lineno, tensor, const_int, const_float, const_str, var_int)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;129m@tvm\u001b[39m\u001b[38;5;241m.\u001b[39mregister_func(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtesting.relax.frontend.nn.test_debug_func\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_debug\u001b[39m(  \u001b[38;5;66;03m# pylint: disable=too-many-arguments\u001b[39;00m\n\u001b[1;32m      3\u001b[0m     lineno: \u001b[38;5;28mstr\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m      8\u001b[0m     var_int: \u001b[38;5;28mint\u001b[39m,\n\u001b[1;32m      9\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m---> 10\u001b[0m     \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtest_frontend_nn_debug.py\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m lineno\n\u001b[1;32m     11\u001b[0m     \u001b[38;5;28;01massert\u001b[39;00m tensor\u001b[38;5;241m.\u001b[39mshape \u001b[38;5;241m==\u001b[39m (\u001b[38;5;241m10\u001b[39m, \u001b[38;5;241m5\u001b[39m)\n\u001b[1;32m     12\u001b[0m     \u001b[38;5;28;01massert\u001b[39;00m const_int \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m1\u001b[39m\n",
      "\u001b[0;31mAssertionError\u001b[0m: "
     ]
    }
   ],
   "source": [
    "@tvm.register_func(\"testing.relax.frontend.nn.test_debug_func\")\n",
    "def _debug(  # pylint: disable=too-many-arguments\n",
    "    lineno: str,\n",
    "    tensor: NDArray,\n",
    "    const_int: int,\n",
    "    const_float: float,\n",
    "    const_str: str,\n",
    "    var_int: int,\n",
    ") -> None:\n",
    "    assert \"test_frontend_nn_debug.py\" in lineno\n",
    "    assert tensor.shape == (10, 5)\n",
    "    assert const_int == 1\n",
    "    assert const_float == 2.0\n",
    "    assert const_str == \"test\"\n",
    "    assert var_int == 8\n",
    "\n",
    "class Layer(nn.Module):\n",
    "    def forward(self, x: nn.Tensor, v: tir.Var):  # pylint: disable=invalid-name\n",
    "        op.debug_func(\"testing.relax.frontend.nn.test_debug_func\", x, 1, 2.0, \"test\", v)\n",
    "        return x\n",
    "\n",
    "model = Layer().jit(\n",
    "    spec={\n",
    "        \"forward\": {\n",
    "            \"x\": spec.Tensor([10, 5], dtype=\"float32\"),\n",
    "            \"v\": \"int\",\n",
    "        },\n",
    "    },\n",
    "    debug=True,\n",
    ")\n",
    "x = torch.rand((10, 5), dtype=torch.float32)  # pylint: disable=invalid-name\n",
    "y = model[\"forward\"](x, 8)  # pylint: disable=invalid-name\n",
    "assert isinstance(y, torch.Tensor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0ae774c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
